#pragma once

#include <vector>
#include <memory>
#include <string>
#include <future>
#include <opencv2/opencv.hpp>

#include "object_detector.hpp"
#include "trt_engine/trt_tensor.h"
#include "yolo_define.h"

_YOLO_NAMESPACE_BEGINE

enum class Type : int
{
    V5 = 0,
    X = 1,
    V3 = V5
};

enum class NMSMethod : int
{
    CPU = 0,    // General, for estimate mAP
    FastGPU = 1 // Fast NMS with a small loss of accuracy in corner cases
};

void image_to_tensor(const cv::Mat &image, std::shared_ptr<Tensor> &tensor, Type type, int ibatch);

class __declspec(dllexport) Infer
{
public:
    virtual std::shared_future<BoxArray> commit(const cv::Mat &image) = 0;
    virtual std::vector<std::shared_future<BoxArray>> commits(const std::vector<cv::Mat> &images) = 0;
};

std::shared_ptr<Infer> __declspec(dllexport) create_infer(
    const std::string &engine_file, Type type, int gpuid,
    float confidence_threshold = 0.25f, float nms_threshold = 0.5f,
    NMSMethod nms_method = NMSMethod::FastGPU, int max_objects = 1024,
    bool use_multi_preprocess_stream = false);

const char *type_name(Type type);

_YOLO_NAMESPACE_END
